#include "algorithm.h"
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <assert.h>
#include "direction.h"
#include "cbuf.h"


#include <stdio.h>

#define WIDTH (space.xend - space.xbegin)
#define HEIGHT (space.yend - space.ybegin)

void* algorithm_memblob = NULL;
int algorithm_memsize = 0;

typedef struct
{
	Direction prev_dir;  // points to where we came from
	int cost;            // cost of path so far
	int cost_estimate;   // total estimated cost
	Coord cursor;
	int length;          // length of path so far
	int length_estimate; // estimated length (for early abort)
} AstarCandidate;

typedef struct
{
	Coord coord;
	float distance;
} FloodFillCandidate;

static struct
{
	AstarCandidate* path_stack;
	int path_stack_size;
	int max_path_size; // max length of any single path
	Direction* prev; // [width*height*depth] Back-pointers for A*
	Coord* solution;
	int solution_size;
} astar_memory;

static struct
{
	Cbuf queue;        // Queue of search coords
	bool* visited;     // [width/xstep*height/ystep] Visited-marker data
	int visited_size;  // # of bytes
	int visited_witdh; // width of the grid, i.e. width/xstep.
} flood_search_memory;

static struct
{
	FloodFillCandidate* queue;
	int capacity;
	int head;
	int tail;
	float* visited; // [width*height] Visited-distance data
} flood_fill_memory;

static void astar_prepare (Space space);
static void flood_search_prepare (Space space, int stepx, int stepy);
static void flood_fill_prepare (Rect space);

void algorithm_init (int memsize)
{
	assert(algorithm_memblob == NULL);
	algorithm_memblob = malloc (memsize);
	algorithm_memsize = memsize;
	
	cbuf_init (&flood_search_memory.queue, sizeof(Coord), 10);
}

void algorithm_cleanup ()
{
	free (algorithm_memblob);
	algorithm_memblob = NULL;
	algorithm_memsize = 0;
	
	cbuf_cleanup (&flood_search_memory.queue);
}

void astar_search (Coord from, Coord to, Space space, AstarResult* result,
                   bool (*can_go) (Coord from, Coord to),
                   int (*cost_step) (Coord pprev, Coord prev, Coord next),
                   int (*cost_estimate) (Coord coord),
                   bool (*visit_handler) (Coord coord))
{
	astar_prepare (space);
	
	#define IDX(X,Y,Z) (((Z)-space.zbegin)*WIDTH*HEIGHT + ((Y)-space.ybegin)*WIDTH + ((X)-space.xbegin))
	#define PREV(C) astar_memory.prev[IDX((C).x,(C).y,(C).z)]
	#define PREVC(C) direction_apply((C), PREV(C), 1);
	
	result->path = astar_memory.solution;
	result->found = false;
	result->overflow = false;
	
	astar_memory.path_stack[0] = (AstarCandidate) { DIR_NONE, 0, 0, from, 0, 0 };
	int stack_size = 1;
	
	int simple_dist (Coord c)
	{
		return abs(c.x - to.x) + abs(c.y - to.y) + abs(c.z - to.z);
	}
	if (cost_estimate == NULL) cost_estimate = &simple_dist;
	
	AstarCandidate pop ()
	{
		assert(stack_size > 0);
		return astar_memory.path_stack[--stack_size];
	}
	
	void push (AstarCandidate candidate) 
	{ 
		if (stack_size >= astar_memory.path_stack_size) {
			// make room by violently pushing a number of the less
			// favorable path candidates off the search stack
			int discard_chunk = ASTAR_MIN_PATHS / 5;
			int keep_chunk = astar_memory.path_stack_size - discard_chunk;
			memmove(astar_memory.path_stack,
			        &astar_memory.path_stack[discard_chunk],
			        keep_chunk * sizeof(AstarCandidate));
			stack_size = keep_chunk;
			result->overflow = true;
		}
		
		int insert = stack_size; // index to push the candidate (respect priority)
		while (insert > 0) {
			AstarCandidate* next = &astar_memory.path_stack[insert-1];
			if (candidate.cost_estimate <= next->cost_estimate) {
				break;
			} else {
				insert--;
			}
		}

		if (insert < stack_size) {
			memmove(&astar_memory.path_stack[insert+1],
			        &astar_memory.path_stack[insert],
			        (stack_size-insert) * sizeof(AstarCandidate));
		}
		
		astar_memory.path_stack[insert] = candidate;
		stack_size++;
	}
	
	void consider (Coord next, int cost_before, int length, Direction prev_dir)
	{
		if ((next.x < space.xbegin) || (next.x >= space.xend) ||
		    (next.y < space.ybegin) || (next.y >= space.yend) ||
		    (next.z < space.zbegin) || (next.z >= space.zend)) 
		{
			return; // out of bounds
		}
		
		if (COORDS_EQUAL(next, from) || (PREV(next) != DIR_NONE)) return; // already been here
		
		Coord prevc = direction_apply(next, prev_dir, 1);
		Coord pprevc = PREVC(prevc);
		
		if (!can_go(prevc, next)) return; // blocked
		if (length > astar_memory.max_path_size) return; // path too long
		
		AstarCandidate candidate;
		
		candidate.prev_dir = prev_dir;
		candidate.cursor = next;
		candidate.cost = cost_before + cost_step (pprevc, prevc, next);
		candidate.cost_estimate = candidate.cost + cost_estimate (next);
		candidate.length = length;
		candidate.length_estimate = length + simple_dist (next);
		push (candidate);
	}
	
	while (stack_size > 0) {
		AstarCandidate candidate = pop();
		Coord c = candidate.cursor;
		
		if (PREV(c) != DIR_NONE) continue; // already been here
		
		PREV(c) = candidate.prev_dir;
		
		if (((visit_handler != NULL) && visit_handler(c)) || COORDS_EQUAL(c, to))
		{
			result->found = true;
			result->length = 1;
			result->path[0] = c;
			
			Direction prev;
			while ((prev = PREV(c)) != DIR_NONE) {
				assert(result->length < astar_memory.solution_size);
				c = direction_apply (c, prev, 1);
				result->path[result->length++] = c;
			}
			
			return;
		}
		
		int cost = candidate.cost;
		int length = candidate.length + 1;
		consider ((Coord) { c.x-1, c.y, c.z }, cost, length, DIR_RIGHT);
		consider ((Coord) { c.x+1, c.y, c.z }, cost, length, DIR_LEFT);
		consider ((Coord) { c.x, c.y-1, c.z }, cost, length, DIR_DOWN);
		consider ((Coord) { c.x, c.y+1, c.z }, cost, length, DIR_UP);
		consider ((Coord) { c.x, c.y, c.z-1 }, cost, length, DIR_LEVELUP);
		consider ((Coord) { c.x, c.y, c.z+1 }, cost, length, DIR_LEVELDOWN);
	}
	
	result->found = false;
	result->length = 0;
	result->path = NULL;
	
	#undef IDX
	#undef PREV
}

bool flood_search (Coord start, Space space, int stepx, int stepy,
                   bool (*can_go) (Coord from, Coord to),
                   bool (*visit_handler) (Coord coord))
{
	flood_search_prepare (space, stepx, stepy);
	
	#define IDX(X,Y) (assert(((X)>=space.xbegin)&&((X)<space.xend)&&((Y)>=space.ybegin)&&((Y)<space.yend)), ((Y)-space.ybegin)/stepy*flood_search_memory.visited_witdh + ((X)-space.xbegin)/stepx)
	#define MARK(X,Y) flood_search_memory.visited[IDX(X,Y)]
	#define QUEUE flood_search_memory.queue
	
	bool layer_search (int z) {
		memset (flood_search_memory.visited, 0, flood_search_memory.visited_size);
		cbuf_clear (&QUEUE);
		start.z = z;
		cbuf_push (&QUEUE, &start);

		while (QUEUE.count > 0) {
			Coord search = *CBUF_POP(&QUEUE,Coord);
			bool visit = visit_handler (search);
			
			if (visit) return true;
			
			void consider (int x, int y)
			{
				if ((x < space.xbegin) || (x >= space.xend) ||
				    (y < space.ybegin) || (y >= space.yend))
				{
					return; // out of bounds
				}
				Coord c = { x, y, z };
				if (!MARK(x,y)) {
					MARK(x,y) = true;
					if (can_go(search, c)) {
						cbuf_push (&QUEUE, &c);
					}
				}
			}
			
			consider (search.x - stepx, search.y);
			consider (search.x + stepx, search.y);
			consider (search.x, search.y - stepy);
			consider (search.x, search.y + stepy);
			consider (search.x - stepx, search.y - stepy);
			consider (search.x + stepx, search.y - stepy);
			consider (search.x - stepx, search.y + stepy);
			consider (search.x + stepx, search.y + stepy);
		}
		
		return false;
	}
	
	int upper = start.z;
	int lower = start.z - 1;
	
	while ((upper < space.zend) || (lower >= space.zbegin)) {
		if (upper < space.zend) {
			if (layer_search (upper++)) return true;
		}
		if (lower >= space.zbegin) {
			if (layer_search (lower--)) return true;
		}
	}
	
	return false;
	
	#undef IDX
	#undef MARK
	#undef QUEUE
}

void flood_fill (Coord start, Rect space, float max_distance,
                 bool (*can_go) (Coord from, Coord to),
                 void (*visit_handler) (Coord coord, float distance))
{
	flood_fill_prepare (space);
	
	#define IDX(X,Y) (((Y)-space.top)*(space.right - space.left) + ((X)-space.left))
	#define MARK(X,Y) flood_fill_memory.visited[IDX((X),(Y))]
	#define QUEUE flood_fill_memory.queue
	#define TAIL flood_fill_memory.tail
	#define HEAD flood_fill_memory.head
	#define CAPACITY flood_fill_memory.capacity
	
	TAIL = 0;
	HEAD = 1;
	QUEUE[0] = (FloodFillCandidate) { start, 0 };
	
	while (TAIL != HEAD) {
		FloodFillCandidate candidate = QUEUE[TAIL++ % CAPACITY];

		void consider (int x, int y, float d)
		{
			if ((x < space.left) || (x >= space.right) ||
				(y < space.top) || (y >= space.bottom))
			{
				return; // out of bounds
			}
			Coord c = { x, y, space.z };
			if ((d <= max_distance) && can_go (candidate.coord, c) && (MARK(x,y) > d)) {
				MARK(x,y) = d;
				assert((HEAD-TAIL) < CAPACITY);
				int idx = HEAD++ % CAPACITY;
				QUEUE[idx] = (FloodFillCandidate) { c, d };
			}
		}

		consider (candidate.coord.x - 1, candidate.coord.y, candidate.distance + 1.0);
		consider (candidate.coord.x + 1, candidate.coord.y, candidate.distance + 1.0);
		consider (candidate.coord.x, candidate.coord.y - 1, candidate.distance + 1.0);
		consider (candidate.coord.x, candidate.coord.y + 1, candidate.distance + 1.0);
		consider (candidate.coord.x - 1, candidate.coord.y - 1, candidate.distance + M_SQRT2);
		consider (candidate.coord.x + 1, candidate.coord.y - 1, candidate.distance + M_SQRT2);
		consider (candidate.coord.x - 1, candidate.coord.y + 1, candidate.distance + M_SQRT2);
		consider (candidate.coord.x + 1, candidate.coord.y + 1, candidate.distance + M_SQRT2);
	}
	
	for(int y = space.top; y < space.bottom; y++)
	for(int x = space.left; x < space.right; x++) {
		float distance = MARK(x,y);
		if (isfinite(distance)) {
			visit_handler ((Coord) { x, y, space.z }, distance);
		}
	}
	
	#undef IDX
	#undef MARK
	#undef QUEUE
	#undef TAIL
	#undef HEAD
	#undef CAPACITY
}

// Partitions the algorithm_memblob for use in astar_search
// [ path_stack | prev pointers ]
static void astar_prepare (Space space)
{
	assert(algorithm_memblob != NULL);
	memset(algorithm_memblob, 0, algorithm_memsize);

	astar_memory.max_path_size = ASTAR_PATH_LENGTH;
	
	astar_memory.prev = algorithm_memblob;
	int prev_size = (space.xend - space.xbegin) * 
	                (space.yend - space.ybegin) * 
	                (space.zend - space.zbegin) * sizeof(Direction);
	
	astar_memory.path_stack_size = (algorithm_memsize - prev_size) / sizeof(AstarCandidate);
	astar_memory.path_stack = (AstarCandidate*) (algorithm_memblob + prev_size);
	
	astar_memory.solution_size = (algorithm_memsize - prev_size) / sizeof(Coord);
	astar_memory.solution = (Coord*) (algorithm_memblob + prev_size);
	
	assert(astar_memory.path_stack_size >= ASTAR_MIN_PATHS);
	assert(astar_memory.solution_size >= ASTAR_PATH_LENGTH);
}

// Partitions the algorithm_memblob for use in flood_search
// [ visited | queue ]
static void flood_search_prepare (Space space, int stepx, int stepy)
{
	assert(algorithm_memblob != NULL);
	
	flood_search_memory.visited = (bool*) algorithm_memblob;
	flood_search_memory.visited_witdh = ((space.xend - space.xbegin - 1) / stepx + 1);
	flood_search_memory.visited_size =
		(flood_search_memory.visited_witdh * 
	   ((space.yend - space.ybegin - 1) / stepy + 1)) *
	   sizeof(bool);
	memset (flood_search_memory.visited, 0, flood_search_memory.visited_size);
	
	assert(flood_search_memory.visited_size <= algorithm_memsize);
}

// Partitions the algorithm_memblob for use in flood_search
// [ visited | queue ]
static void flood_fill_prepare (Rect space)
{
	assert(algorithm_memblob != NULL);
	
	flood_fill_memory.visited = (float*) algorithm_memblob;
	int visited_len = (space.right - space.left) * (space.bottom - space.top);
	int visited_size = visited_len * sizeof(float);
	for (int i = 0; i < visited_len; i++) {
		flood_fill_memory.visited[i] = INFINITY;
	}
	
	flood_fill_memory.capacity = (algorithm_memsize - visited_size) / sizeof(FloodFillCandidate);
	flood_fill_memory.queue = (FloodFillCandidate*) (algorithm_memblob + visited_size);
	flood_fill_memory.head = 0;
	flood_fill_memory.tail = 0;
	
	assert(flood_fill_memory.capacity >= FLOOD_MIN_COORDS);
}
